#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
import argparse
import os
import sys
from time import sleep
import datetime

import grpc

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

SWITCH_TO_HOST_PORT = 1

# 计数器日志文件路径
COUNTER_LOG_FILE = "logs/tunnel_counters.log"
DIRECTION_LOG_DIR = "logs/direction_logs"

def ensure_log_dirs():
    """确保日志目录存在"""
    os.makedirs(DIRECTION_LOG_DIR, exist_ok=True)
    os.makedirs(os.path.dirname(COUNTER_LOG_FILE), exist_ok=True)

def save_counters_to_file(counter_data, filename=COUNTER_LOG_FILE):
    """将计数器数据保存到主日志文件"""
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(filename, "a") as f:
        f.write(f"\n----- {timestamp} -----\n")
        for line in counter_data:
            f.write(line + "\n")
        f.write("-" * 50 + "\n")

def save_direction_counters(direction_name, counter_data):
    """将每个方向的计数器保存到单独的文件"""
    filename = os.path.join(DIRECTION_LOG_DIR, f"{direction_name.replace(' ', '_').replace('->', 'to')}.log")
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(filename, "a") as f:
        f.write(f"{timestamp} - ")
        for line in counter_data:
            f.write(line + " | ")
        f.write("\n")

def writeIpv4_lpmRules(p4info_helper, sw, dst_ip_addr, tunnel_id):
    """
    Installs tunnel ingress rule in the ipv4_lpm table
    """
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    sw.WriteTableEntry(table_entry)
    print("Installed ingress tunnel rule on %s for tunnel %d" % (sw.name, tunnel_id))

def writeMyTunnel_exactRules(p4info_helper, ingress_sw, egress_sw, tunnel_id, dst_eth_addr):
    """
    Installs transit and egress rules for tunnel forwarding
    """
    # Determine switch-to-switch port based on tunnel ID
    SWITCH_TO_SWITCH_PORT = None
    if tunnel_id == 100:  # s1 -> s2
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 101:  # s2 -> s1
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 200:  # s1 -> s3
        SWITCH_TO_SWITCH_PORT = 3
    elif tunnel_id == 201:  # s3 -> s1
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 300:  # s2 -> s3
        SWITCH_TO_SWITCH_PORT = 3
    elif tunnel_id == 301:  # s3 -> s2
        SWITCH_TO_SWITCH_PORT = 3

    # Transit rule on ingress switch
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": SWITCH_TO_SWITCH_PORT
        })
    ingress_sw.WriteTableEntry(table_entry)
    print("Installed transit tunnel rule on %s for tunnel %d" % (ingress_sw.name, tunnel_id))

    # Egress rule on egress switch
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)
    print("Installed egress tunnel rule on %s for tunnel %d" % (egress_sw.name, tunnel_id))

def writeArp_matchRules(p4info_helper, sw, match_fields_1, match_fields_2, action_params):
    """
    Installs ARP reply rules
    """
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.arp_match",
        match_fields={
            "hdr.arp.oper": match_fields_1,
            "hdr.arp.tpa": match_fields_2
        },
        action_name="MyIngress.send_arp_reply",
        action_params=action_params
    )
    sw.WriteTableEntry(table_entry)
    print("Installed ARP rule on %s" % sw.name)

def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def get_counter_value(p4info_helper, sw, counter_name, index):
    """获取计数器的值并返回格式化字符串"""
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            return "%s %s %d: %d packets (%d bytes)" % (
                sw.name, counter_name, index,
                counter.data.packet_count, counter.data.byte_count
            )
    return "%s %s %d: 0 packets (0 bytes)" % (sw.name, counter_name, index)

def print_and_save_counters(p4info_helper, s1, s2, s3):
    """读取、打印并保存所有计数器"""
    print('\n----- Reading tunnel counters -----')

    all_counter_data = []
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # s1 -> s2
    print('\n--- s1 -> s2 ---')
    s1_ingress_100 = get_counter_value(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
    s2_egress_100 = get_counter_value(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
    print(s1_ingress_100)
    print(s2_egress_100)
    all_counter_data.extend([s1_ingress_100, s2_egress_100])
    save_direction_counters("s1_to_s2", [s1_ingress_100, s2_egress_100])

    # s2 -> s1
    print('\n--- s2 -> s1 ---')
    s2_ingress_101 = get_counter_value(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
    s1_egress_101 = get_counter_value(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)
    print(s2_ingress_101)
    print(s1_egress_101)
    all_counter_data.extend([s2_ingress_101, s1_egress_101])
    save_direction_counters("s2_to_s1", [s2_ingress_101, s1_egress_101])

    # s1 -> s3
    print('\n--- s1 -> s3 ---')
    s1_ingress_200 = get_counter_value(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
    s3_egress_200 = get_counter_value(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
    print(s1_ingress_200)
    print(s3_egress_200)
    all_counter_data.extend([s1_ingress_200, s3_egress_200])
    save_direction_counters("s1_to_s3", [s1_ingress_200, s3_egress_200])

    # s3 -> s1
    print('\n--- s3 -> s1 ---')
    s3_ingress_201 = get_counter_value(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
    s1_egress_201 = get_counter_value(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)
    print(s3_ingress_201)
    print(s1_egress_201)
    all_counter_data.extend([s3_ingress_201, s1_egress_201])
    save_direction_counters("s3_to_s1", [s3_ingress_201, s1_egress_201])

    # s2 -> s3
    print('\n--- s2 -> s3 ---')
    s2_ingress_300 = get_counter_value(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
    s3_egress_300 = get_counter_value(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
    print(s2_ingress_300)
    print(s3_egress_300)
    all_counter_data.extend([s2_ingress_300, s3_egress_300])
    save_direction_counters("s2_to_s3", [s2_ingress_300, s3_egress_300])

    # s3 -> s2
    print('\n--- s3 -> s2 ---')
    s3_ingress_301 = get_counter_value(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
    s2_egress_301 = get_counter_value(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)
    print(s3_ingress_301)
    print(s2_egress_301)
    all_counter_data.extend([s3_ingress_301, s2_egress_301])
    save_direction_counters("s3_to_s2", [s3_ingress_301, s2_egress_301])

    print()  # 空行分隔

    # 保存到主日志文件
    save_counters_to_file(all_counter_data)

    return all_counter_data

def printGrpcError(e):
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print("(%s)" % status_code.name, end=' ')
    traceback = sys.exc_info()[2]
    print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))

def main(p4info_file_path, bmv2_file_path):
    # 确保日志目录存在
    ensure_log_dirs()

    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connection objects for s1, s2 and s3
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message to establish this controller as master
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on the switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        # Write ipv4_lpm table entries for all paths
        # s1 rules
        writeIpv4_lpmRules(p4info_helper, sw=s1, dst_ip_addr="10.0.2.2", tunnel_id=100)
        writeIpv4_lpmRules(p4info_helper, sw=s1, dst_ip_addr="10.0.3.3", tunnel_id=200)

        # s2 rules
        writeIpv4_lpmRules(p4info_helper, sw=s2, dst_ip_addr="10.0.1.1", tunnel_id=101)
        writeIpv4_lpmRules(p4info_helper, sw=s2, dst_ip_addr="10.0.3.3", tunnel_id=300)

        # s3 rules
        writeIpv4_lpmRules(p4info_helper, sw=s3, dst_ip_addr="10.0.1.1", tunnel_id=201)
        writeIpv4_lpmRules(p4info_helper, sw=s3, dst_ip_addr="10.0.2.2", tunnel_id=301)

        # Write myTunnel_exact table entries for all paths
        writeMyTunnel_exactRules(p4info_helper, ingress_sw=s1, egress_sw=s2, tunnel_id=100, dst_eth_addr="08:00:00:00:02:22")
        writeMyTunnel_exactRules(p4info_helper, ingress_sw=s1, egress_sw=s3, tunnel_id=200, dst_eth_addr="08:00:00:00:03:33")
        writeMyTunnel_exactRules(p4info_helper, ingress_sw=s2, egress_sw=s1, tunnel_id=101, dst_eth_addr="08:00:00:00:01:11")
        writeMyTunnel_exactRules(p4info_helper, ingress_sw=s2, egress_sw=s3, tunnel_id=300, dst_eth_addr="08:00:00:00:03:33")
        writeMyTunnel_exactRules(p4info_helper, ingress_sw=s3, egress_sw=s1, tunnel_id=201, dst_eth_addr="08:00:00:00:01:11")
        writeMyTunnel_exactRules(p4info_helper, ingress_sw=s3, egress_sw=s2, tunnel_id=301, dst_eth_addr="08:00:00:00:02:22")

        # Write ARP table entries
        writeArp_matchRules(p4info_helper, s1, [1], ["10.0.1.10", 32], {"macAddr": "08:00:00:00:01:00"})
        writeArp_matchRules(p4info_helper, s2, [1], ["10.0.2.20", 32], {"macAddr": "08:00:00:00:02:00"})
        writeArp_matchRules(p4info_helper, s3, [1], ["10.0.3.30", 32], {"macAddr": "08:00:00:00:03:00"})

        # Read table entries from all switches
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        print("\n=== 开始记录计数器数据 ===")
        print(f"主日志文件: {COUNTER_LOG_FILE}")
        print(f"方向日志目录: {DIRECTION_LOG_DIR}")
        print("=" * 50)

        # Print the tunnel counters every 2 seconds and save to files
        while True:
            sleep(2)
            print_and_save_counters(p4info_helper, s1, s2, s3)

    except KeyboardInterrupt:
        print("\n=== 停止记录计数器数据 ===")
        print("计数器数据已保存到以下文件：")
        print(f"- 主日志: {COUNTER_LOG_FILE}")
        print(f"- 各方向日志: {DIRECTION_LOG_DIR}/")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
